import random
import pickle
import os
import math

def generate_prime():
    return _dec_to_bin(_random_prime())

def pick_g_and_x(p):
    p = int(p, 2)
    g = random.randint(2, p-1)
    x = random.randint(2, p-1)
    return _dec_to_bin(g), _dec_to_bin(x)

def calculate_y(p, g, x):
    p = int(p, 2)
    g = int(g, 2)
    x = int(x, 2)
    y = (g**x) % p
    return _dec_to_bin(y)

def get_keys(p, g, x, y):
    public = y + g + p
    private = x
    return public, private

def _random_prime():
    """
    Generates a prime number under 2**16. Uses pre-generated list of primes.
    """
    primes = pickle.loads(__loader__.get_data("primes.pkl"))
    while primes[0] < 2**8:
        primes.pop(0)
    return random.choice(primes)

def _dec_to_bin(decimal):
    binary_reps = {'0': '000',
                   '1': '001',
                   '2': '010',
                   '3': '011',
                   '4': '100',
                   '5': '101',
                   '6': '110',
                   '7': '111'
    }
    oct_rep = str(oct(decimal))
    binary_string = ""
    for c in oct_rep:
        if c != "L":
            binary_string += binary_reps[c]
    if len(binary_string) > 16:
        binary_string = binary_string[len(binary_string) - 16:]
    while len(binary_string) < 16:
        binary_string = "0" + binary_string
    return binary_string